import pandas as pd
import os


class GenericExcel:
    def __init__(self, file_path, types=None):
        """
        初始化 Excel 处理类

        :param file_path: Excel 文件路径
        :param types: 字典，指定列的数据类型（如 {'color': 'category', 'size': 'category'}）
        """
        self.file_path = file_path
        self.types = types or {}
        self.metadata_sheet = "__metadata__"

        # 检查文件是否存在
        if os.path.exists(file_path):
            # 读取数据和元数据
            self._read_with_metadata()
        else:
            # 创建新 DataFrame
            self.df = pd.DataFrame()
            print(f"创建新文件: {file_path}")

    def _read_with_metadata(self):
        """读取 Excel 文件及其元数据"""
        try:
            # 尝试读取元数据
            metadata = pd.read_excel(self.file_path, sheet_name=self.metadata_sheet)
            # 从元数据中提取类型信息
            self.types = dict(zip(metadata['column'], metadata['dtype']))
            print("从元数据加载类型信息")
        except:
            print("未找到元数据，使用提供的类型信息")

        # 读取主数据
        self.df = pd.read_excel(self.file_path)

        # 应用类型转换
        self._apply_types()

    def _apply_types(self):
        """应用类型转换到 DataFrame"""
        for col, dtype in self.types.items():
            if col in self.df.columns:
                if dtype == 'category':
                    # 获取所有唯一值作为类别
                    unique_vals = self.df[col].dropna().unique().tolist()
                    self.df[col] = pd.Categorical(
                        self.df[col],
                        categories=unique_vals
                    )
                else:
                    try:
                        self.df[col] = self.df[col].astype(dtype)
                    except Exception as e:
                        print(f"无法将列 '{col}' 转换为 {dtype}: {e}")

    def _save_metadata(self):
        """保存元数据到单独的 sheet"""
        metadata = pd.DataFrame({
            'column': list(self.types.keys()),
            'dtype': list(self.types.values())
        })
        return metadata

    def sub_rows(self, filter_func):
        """获取满足条件的行子集"""
        return self.df[self.df.apply(filter_func, axis=1)]

    def drop_rows(self, filter_func):
        """删除满足条件的行"""
        self.df = self.df[~self.df.apply(filter_func, axis=1)]

    def append_row(self, row):
        """添加新行到 DataFrame"""
        # 处理分类列的新值
        for col, dtype in self.types.items():
            if dtype == 'category' and col in row:
                # 确保该列存在
                if col not in self.df.columns:
                    self.df[col] = pd.Categorical([], categories=[])

                # 如果新值不在现有类别中，扩展类别
                if row[col] not in self.df[col].cat.categories:
                    self.df[col] = self.df[col].cat.add_categories([row[col]])

        # 添加新行
        new_row = pd.DataFrame([row])
        self.df = pd.concat([self.df, new_row], ignore_index=True)

    def save(self, output_path=None):
        """保存到 Excel 文件，包含元数据"""
        save_path = output_path or self.file_path

        # 保存主数据和元数据
        with pd.ExcelWriter(save_path) as writer:
            self.df.to_excel(writer, sheet_name='Data', index=False)

            # 保存元数据
            if self.types:
                metadata = self._save_metadata()
                metadata.to_excel(writer, sheet_name=self.metadata_sheet, index=False)

    def get_dataframe(self):
        """获取当前 DataFrame"""
        return self.df.copy()

    def set_dataframe(self, df):
        """设置新的 DataFrame"""
        self.df = df.copy()
        self._apply_types()

    def update_types(self, new_types):
        """更新类型信息"""
        self.types.update(new_types)
        self._apply_types()

    def add_column(self, col_name, dtype, default_value=None):
        """添加新列并指定类型"""
        # 初始化新列
        self.df[col_name] = default_value

        # 设置类型
        self.types[col_name] = dtype
        self._apply_types()
